# -*- coding:utf-8 -*-
from wmcloud_dataclient import Client
from cassandra.cluster import Cluster
from cassandra import ConsistencyLevel
from cassandra.query import BatchStatement, BatchType
from datetime import *
import traceback
import json
import time
import os

def capture_wmcloud_stock_raw_data(stock_ticker_list = [], try_count = 3):
    client = Client('93ef5ebba60c5f1da190801750612423db818daad66155b6b191417b19d331b9')
    cluster = Cluster(protocol_version = 4)
    session = cluster.connect()
    session.set_keyspace('wmcloud')
    insert_sock = session.prepare('insert into wmcloud.sz_sh_stock_daily_trade_info(ticker, tradeDate, secID, secShortName, exchangeCD, preClosePrice, actPreClosePrice, openPrice, highestPrice, lowestPrice, closePrice, turnoverVol, turnoverValue, dealAmount, turnoverRate, accumAdjFactor, negMarketValue, marketValue) values(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)')
    end_datetime = datetime.today()
    weekend = int(end_datetime.strftime('%u')) - 5
    if weekend > 0:
         end_datetime = end_datetime - timedelta(days=weekend)
    end_time = str(end_datetime.strftime('%Y%m%d'))
    for stock_ticker in stock_ticker_list:
        row = session.execute("select tradedate from wmcloud.sz_sh_stock_daily_trade_info where ticker='%s' order by tradedate desc limit 1"%(stock_ticker))
        if len(row) > 0:
            if row[0].tradedate.strftime('%Y%m%d') == end_time:
                # 说明已经下载了最新数据了不需要进一步的下载
                continue
            sdatet = row[0].tradedate + timedelta(days = 1)
            start_time = sdatet.strftime('%Y%m%d')
        else:
            start_time = '19970101'
        # 首先还是要通过cassandra的查询来获取某个股票的所保存的最近日期
        url = '/api/market/getMktEqud.json?field=secID,tradeDate,ticker,secShortName,exchangeCD,preClosePrice,actPreClosePrice,openPrice,highestPrice,lowestPrice,closePrice,tradeDate,turnoverVol,turnoverValue,dealAmount,turnoverRate,accumAdjFactor,negMarketValue,marketValue&beginDate=%s&endDate=%s&ticker=%s'%(start_time, end_time, stock_ticker)
        for try_c in range(try_count):
            try:
                code, result = client.getData(url)
                if code == 200:
                    stock_json = json.loads(result)
                    print(stock_ticker)
                    if 'retCode' not in stock_json:
                        print('error json : ',stock_json)
                    elif stock_json['retCode'] == 1:
                        batch_stmt = BatchStatement(batch_type = BatchType.UNLOGGED, consistency_level = ConsistencyLevel.QUORUM)
                        for inx, data in enumerate(stock_json['data']):
                            # Save it to cassandra (scylladb)
                            if 'accumAdjFactor' not in data:
                                data['accumAdjFactor'] = 0
                            # 这里主要是把时间调整到下午4点这个时候都收盘了
                            trade_date_timestamp = datetime.strptime(data['tradeDate'], '%Y-%m-%d') + timedelta(hours = 16)
                            batch_stmt.add(insert_sock, (data['ticker'], trade_date_timestamp, data['secID'], data['secShortName'], data['exchangeCD'], data['preClosePrice'], data['actPreClosePrice'], data['openPrice'], data['highestPrice'], data['lowestPrice'], data['closePrice'], data['turnoverVol'], data['turnoverValue'], data['dealAmount'] if 'dealAmount' in data else 0, data['turnoverRate'], data['accumAdjFactor'], data['negMarketValue'], data['marketValue']))
                            if inx != 0 and inx % 50 == 0:
                                session.execute(batch_stmt)
                                batch_stmt = BatchStatement(batch_type = BatchType.UNLOGGED, consistency_level = ConsistencyLevel.QUORUM)
                        else:
                            session.execute(batch_stmt)
                    break
            except Exception as e:
                print(e)
                continue
    session.shutdown()
    cluster.shutdown()

def get_stock_ticker_list():
    cluster = Cluster(protocol_version=4)
    session = cluster.connect()
    session.set_keyspace('wmcloud')
    rows = session.execute('select ticker from wmcloud.sz_sh_stock_basic_info')
    stock_ticker_list = []
    for row in rows:
        stock_ticker_list.append(row.ticker)

    session.shutdown()
    cluster.shutdown()
    return stock_ticker_list


if __name__ == '__main__':
    try:
        stock_ticker_list = get_stock_ticker_list()
        capture_wmcloud_stock_raw_data(stock_ticker_list)
    except Exception as e:
        traceback.print_exc()
